package com.aiwiown.snackmq.cluster.raft.closeure;

import com.aiwiown.snackmq.common.handler.RequestContext;
import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.message.MessageStatus;
import com.aiwiown.snackmq.common.message.MessageType;
import com.aiwiown.snackmq.common.metrics.MessageMetricEmitter;
import com.alipay.sofa.jraft.Closure;
import com.alipay.sofa.jraft.Status;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;

/**
 * Raft 状态机应用闭包，用于在 Raft 日志应用后执行回调。
 * 主要负责根据状态机处理结果或 Raft 状态构建相应的响应消息，并发送给请求方。
 */
@Slf4j
public class SnackMQApplyClosure implements Closure {

    /**
     * 由状态机设置的响应消息，包含详细的业务处理结果。
     */
    @Setter
    private Message response;

    /**
     * 请求消息，用于构建响应消息时引用原始请求的信息。
     */
    private final Message requestMessage;

    /**
     * 请求上下文，用于写入响应。
     */
    private final RequestContext ctx;

    /**
     * 消息度量发射器，用于记录消息相关的度量信息。
     */
    private final MessageMetricEmitter metricEmitter;

    /**
     * 构造方法，初始化成员变量。
     *
     * @param requestMessage 请求消息
     * @param ctx 请求上下文
     * @param metricEmitter 消息度量发射器
     */
    public SnackMQApplyClosure(Message requestMessage, RequestContext ctx, MessageMetricEmitter metricEmitter) {
        this.requestMessage = requestMessage;
        this.ctx = ctx;
        this.metricEmitter = metricEmitter;
    }

    /**
     * 实现 Closure 接口的 run 方法，根据业务处理结果或 Raft 状态执行相应的逻辑。
     *
     * @param status Raft 状态
     */
    @Override
    public void run(Status status) {
        // 优先使用状态机在 onApply 中设置的、包含详细业务结果的 response 对象
        if (this.response != null) {
            ctx.writeResponse(this.response);
            if (this.response.getStatus() == MessageStatus.SUCCESS) {
                log.info("Successfully processed and sent response for messageId: {}", requestMessage.getMessageId());
                // 【修复】调用正确的方法来记录已生产的消息数量
                if (this.metricEmitter != null && requestMessage.getType() == MessageType.PRODUCE_REQUEST) {
                    long messageCount = requestMessage.isBatch() ? requestMessage.getMessages().size() : 1;
                    this.metricEmitter.markMessagesProduced(messageCount);
                }
            } else {
                log.warn("Business logic failed for messageId: {}. Reason: {}", requestMessage.getMessageId(), this.response.getErrorMessage());
            }
        } else {
            // 如果状态机没有设置 response（例如对于 DECLARE_EXCHANGE, BIND_QUEUE 等操作），
            // 则根据 Raft 状态构建一个通用的成功/失败响应。
            if (status.isOk()) {
                Message ackMessage = Message.builder()
                        .type(MessageType.RESPONSE)
                        .messageId(requestMessage.getMessageId())
                        // --- 关键修复：明确设置成功状态 ---
                        .status(MessageStatus.SUCCESS)
                        .build();
                ctx.writeResponse(ackMessage);
                log.info("Successfully replicated (generic ACK) for messageId: {}", requestMessage.getMessageId());
            } else {
                sendErrorResponse(ctx, requestMessage, status.getErrorMsg());
            }
        }
    }

    /**
     * 静态辅助方法，用于发送一个标准的错误响应。
     *
     * @param ctx 请求上下文
     * @param request 请求消息
     * @param errorMessage 错误消息
     */
    public static void sendErrorResponse(RequestContext ctx, Message request, String errorMessage) {
        Message nackMessage = Message.builder()
                .type(MessageType.RESPONSE)
                .messageId(request.getMessageId())
                .status(MessageStatus.FAILED)
                .errorMessage(errorMessage)
                .build();
        ctx.writeResponse(nackMessage);
    }

    /**
     * --- 新增方法：发送重定向响应 ---
     * 当前节点不是 Leader 时，向客户端发送重定向到新 Leader 的响应。
     *
     * @param ctx 请求上下文
     * @param request 请求消息
     * @param redirectAddress 新 Leader 的地址
     */
    public static void sendRedirectResponse(RequestContext ctx, Message request, String redirectAddress) {
        Message nackMessage = Message.builder()
                .type(MessageType.RESPONSE)
                .messageId(request.getMessageId())
                .status(MessageStatus.FAILED)
                .errorMessage("Not the leader, please retry.")
                .property("redirectAddress", redirectAddress) // 在属性中告诉客户端新地址
                .build();
        ctx.writeResponse(nackMessage);
    }

    /**
     * --- 新增方法：发送重定向响应 ---
     * 当前节点不是 Leader 时，向客户端发送重定向到新 Leader 的响应。
     *
     * @param ctx 请求上下文
     * @param request 请求消息
     * @param redirectAddress 新 Leader 的地址
     */
    public static void sendRedirectResponse(RequestContext ctx, Message request, String redirectAddress, String errorMessage) {
        Message nackMessage = Message.builder()
                .type(MessageType.RESPONSE)
                .messageId(request.getMessageId())
                .status(MessageStatus.FAILED)
                .errorMessage(errorMessage)
                .property("redirectAddress", redirectAddress) // 在属性中告诉客户端新地址
                .build();
        ctx.writeResponse(nackMessage);
    }

    /**
     * 静态辅助方法，用于发送一个标准的成功响应。
     * 这对于那些不需要等待 Raft 提交，或者业务逻辑本身就是“成功”（如丢弃无法路由的消息）的场景非常有用。
     *
     * @param ctx 请求上下文
     * @param request 请求消息
     */
    public static void sendSuccessResponse(RequestContext ctx, Message request) {
        Message ackMessage = Message.builder()
                .type(MessageType.RESPONSE)
                .messageId(request.getMessageId())
                .status(MessageStatus.SUCCESS)
                .build();
        ctx.writeResponse(ackMessage);
    }
}
